{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PMSM MPC dq current control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we apply a continuous-control set (CCS) model predictive control (MPC) approach to realize the current control of a permanent magnet synchronous motor (PMSM) in rotor field-oriented coordinates. We achieve this by utilizing the GEKKO toolbox [https://gekko.readthedocs.io/en/latest/].\n",
    "\n",
    "![MPC2](img/mpc_structure.png)\n",
    "\n",
    "![MPC1](img/mpc_scheme.png)\n",
    "\n",
    "With the help of the system model the output variables are calculated. By minimizing a cost function, the optimizer determines the optimal control variables for the following N time steps, the so-called prediction horizon. The first manipulated variable is switched on and then the system outputs are measured. When calculating the manipulated variables, the manipulated variable limits are explicitly taken into account.\n",
    "\n",
    "![Limits](img/voltage_limits.png)\n",
    "The quadratic control error area is used as the cost function. Since the control is performed in the rotating dq-coordinates, the voltage limits, which have the shape of a hexagon, are time-variant. The optimization problem is solved iteratively using an Interior-Point Solver, which approximates the limits using barrier functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation of the required toolboxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q git+https://github.com/BYU-PRISM/GEKKO.git\n",
    "!pip install -q git+https://github.com/upb-lea/gym-electric-motor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "import matplotlib\n",
    "from gekko import GEKKO\n",
    "\n",
    "import gym_electric_motor as gem\n",
    "from gym_electric_motor.envs.motors import ActionType, ControlType, Motor, MotorType\n",
    "from gym_electric_motor.physical_systems import ConstantSpeedLoad\n",
    "from gym_electric_motor.reference_generators import MultipleReferenceGenerator, SwitchedReferenceGenerator, \\\n",
    "        TriangularReferenceGenerator, SinusoidalReferenceGenerator, StepReferenceGenerator\n",
    "from gym_electric_motor.visualization.motor_dashboard import MotorDashboard\n",
    "from gym_electric_motor.physical_systems.solvers import ScipySolveIvpSolver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition of a general controller class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Controller:\n",
    "\n",
    "    @classmethod\n",
    "    def make(cls, controller_type, environment, **controller_kwargs):\n",
    "        assert controller_type in _controllers.keys(), f'Controller {controller_type} unknown'\n",
    "        controller = _controllers[controller_type](environment, **controller_kwargs)\n",
    "        return controller\n",
    "\n",
    "    def control(self, state, reference):\n",
    "        pass\n",
    "\n",
    "    def reset(self):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition of the MPC class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MPC(Controller):\n",
    "    def __init__(self, environment, ph=5, ref_idx_q=0, ref_idx_d=1):\n",
    "        # conversion of the coordinate systems\n",
    "        t32 = environment.physical_system.electrical_motor.t_32\n",
    "        q = environment.physical_system.electrical_motor.q\n",
    "        self._backward_transformation = (lambda quantities, eps: t32(q(quantities[::-1], eps)))\n",
    "    \n",
    "        # indices\n",
    "        self.ref_idx_i_q = ref_idx_q\n",
    "        self.ref_idx_i_d = ref_idx_d\n",
    "        self.i_sd_idx = environment.state_names.index('i_sd')\n",
    "        self.i_sq_idx = environment.state_names.index('i_sq')\n",
    "        self.u_a_idx = environment.state_names.index('u_a')\n",
    "        self.u_b_idx = environment.state_names.index('u_b')\n",
    "        self.u_c_idx = environment.state_names.index('u_c')\n",
    "        self.u_sq_idx = environment.state_names.index('u_sq')\n",
    "        self.u_sd_idx = environment.state_names.index('u_sd')\n",
    "        self.omega_idx = environment.state_names.index('omega')\n",
    "        self.epsilon_idx = environment.state_names.index('epsilon')\n",
    "\n",
    "        # motor parameters\n",
    "        self.tau = environment.physical_system.tau\n",
    "        self.limits = environment.physical_system.limits\n",
    "        self.l_q = environment.physical_system.electrical_motor.motor_parameter['l_q']\n",
    "        self.l_d = environment.physical_system.electrical_motor.motor_parameter['l_d']\n",
    "        self.psi_ = environment.physical_system.electrical_motor.motor_parameter['psi_p']\n",
    "        self.r_s = environment.physical_system.electrical_motor.motor_parameter['r_s']\n",
    "        self.p = environment.physical_system.electrical_motor.motor_parameter['p']\n",
    "        self.ph_ = ph\n",
    "\n",
    "    def control(self, state, reference):\n",
    "        # initialize variables\n",
    "        epsilon_el = state[self.epsilon_idx] * self.limits[self.epsilon_idx]\n",
    "        omega = self.p * state[self.omega_idx] * self.limits[self.omega_idx]\n",
    "\n",
    "        ref_q = []\n",
    "        ref_d = []\n",
    "        eps = []\n",
    "        lim_a_up = []\n",
    "        lim_a_low = []\n",
    "        \n",
    "        for i in range(self.ph_):\n",
    "            ref_q.append(reference[self.ref_idx_i_q] * self.limits[self.i_sq_idx])\n",
    "            ref_d.append(reference[self.ref_idx_i_d] * self.limits[self.i_sd_idx])\n",
    "         \n",
    "            eps.append(epsilon_el + (i-1) * self.tau * omega)\n",
    "            lim_a_up.append(2 * self.limits[self.u_a_idx])\n",
    "            lim_a_low.append(-2 * self.limits[self.u_a_idx])\n",
    "        \n",
    "        m = GEKKO(remote=False)\n",
    "        \n",
    "        # defenition of the prediction Horizon\n",
    "        m.time = np.linspace(self.tau, self.tau * self.ph_, self.ph_)\n",
    "\n",
    "        # defenition of the variables\n",
    "        u_d = m.MV(value=state[self.u_sd_idx] * self.limits[self.u_sd_idx])\n",
    "        u_q = m.MV(value=state[self.u_sq_idx] * self.limits[self.u_sq_idx])\n",
    "        u_d.STATUS = 1\n",
    "        u_q.STATUS = 1\n",
    "\n",
    "        u_a_lim_up = m.Param(value=lim_a_up)\n",
    "        u_a_lim_low = m.Param(value=lim_a_low)\n",
    "        sq3 = math.sqrt(3)\n",
    "\n",
    "        i_d = m.SV(value=state[self.i_sd_idx] * self.limits[self.i_sd_idx], lb=-self.limits[self.i_sd_idx], ub=self.limits[self.i_sd_idx] )\n",
    "        i_q = m.SV(value=state[self.i_sq_idx] * self.limits[self.i_sq_idx], lb=-self.limits[self.i_sq_idx], ub=self.limits[self.i_sq_idx])\n",
    "\n",
    "        epsilon = m.Param(value=eps)\n",
    "        \n",
    "        # reference trajectory\n",
    "        traj_d = m.Param(value=ref_d)\n",
    "        traj_q = m.Param(value=ref_q)\n",
    "        \n",
    "        # defenition of the constants\n",
    "        omega = m.Const(value=omega)\n",
    "        psi = m.Const(value=self.psi_)\n",
    "        rs = m.Const(value=self.r_s)\n",
    "        ld = m.Const(value=self.l_d)\n",
    "        lq = m.Const(value=self.l_q)\n",
    "        \n",
    "        # control error\n",
    "        e_d = m.CV()\n",
    "        e_q = m.CV()\n",
    "        e_d.STATUS = 1\n",
    "        e_q.STATUS = 1\n",
    "        \n",
    "        # solver options\n",
    "        m.options.CV_TYPE = 2\n",
    "        m.options.IMODE = 6\n",
    "        m.options.solver = 3\n",
    "        m.options.WEB = 0\n",
    "        m.options.NODES = 2\n",
    "        \n",
    "        # differential equations\n",
    "        m.Equations([ld * i_d.dt() == u_d - rs * i_d + omega * lq * i_q,\n",
    "                     lq * i_q.dt() == u_q - rs * i_q - omega * ld * i_d - omega * psi])\n",
    "        \n",
    "        # cost function\n",
    "        m.Equations([e_d == (i_d - traj_d), e_q == (i_q - traj_q)])\n",
    "        \n",
    "        # voltage limitations\n",
    "        m.Equation(u_a_lim_up >= 3/2 * m.cos(epsilon) * u_d - 3/2 * m.sin(epsilon) * u_q - sq3/2 * m.sin(epsilon) * u_d - sq3/2 * m.cos(epsilon) * u_q)\n",
    "        m.Equation(u_a_lim_low <= 3 / 2 * m.cos(epsilon) * u_d - 3 / 2 * m.sin(epsilon) * u_q - sq3 / 2 * m.sin(epsilon) * u_d - sq3 / 2 * m.cos(epsilon) * u_q)\n",
    "        m.Equation(u_a_lim_up >= sq3 * m.sin(epsilon) * u_d + sq3 * m.cos(epsilon) * u_q)\n",
    "        m.Equation(u_a_lim_low <= sq3 * m.sin(epsilon) * u_d + sq3 * m.cos(epsilon) * u_q)\n",
    "        m.Equation(u_a_lim_up >= -3 / 2 * m.cos(epsilon) * u_d + 3 / 2 * m.sin(epsilon) * u_q - sq3 / 2 * m.sin(epsilon) * u_d - sq3 / 2 * m.cos(epsilon) * u_q)\n",
    "        m.Equation(u_a_lim_low <= -3 / 2 * m.cos(epsilon) * u_d + 3 / 2 * m.sin(epsilon) * u_q - sq3 / 2 * m.sin(epsilon) * u_d - sq3 / 2 * m.cos(epsilon) * u_q)\n",
    "        \n",
    "        # object to minimize\n",
    "        m.Obj(e_d)\n",
    "        m.Obj(e_q)\n",
    "        \n",
    "        # solving optimization problem\n",
    "        m.solve(disp=False)\n",
    "        \n",
    "        # additional voltage limitation\n",
    "        u_a, u_b, u_c = self._backward_transformation((u_q.NEWVAL, u_d.NEWVAL), epsilon_el)\n",
    "        u_max = max(np.absolute(u_a - u_b), np.absolute(u_b - u_c), np.absolute(u_c - u_a))\n",
    "        if u_max >= 2 * self.limits[self.u_a_idx]:\n",
    "            u_a = u_a / u_max * 2 * self.limits[self.u_a_idx]\n",
    "            u_b = u_b / u_max * 2 * self.limits[self.u_a_idx]\n",
    "            u_c = u_c / u_max * 2 * self.limits[self.u_a_idx]\n",
    "        \n",
    "        # Zero Point Shift\n",
    "        u_0 = 0.5 * (max(u_a, u_b, u_c) + min(u_a, u_b, u_c))\n",
    "        u_a -= u_0\n",
    "        u_b -= u_0\n",
    "        u_c -= u_0\n",
    "        \n",
    "        # normalization of the manipulated variables\n",
    "        u_a /= self.limits[self.u_a_idx]\n",
    "        u_b /= self.limits[self.u_b_idx]\n",
    "        u_c /= self.limits[self.u_c_idx]\n",
    "        \n",
    "        return u_a, u_b, u_c\n",
    "\n",
    "    def reset(self):\n",
    "        None\n",
    "\n",
    "\n",
    "_controllers = {\n",
    "    'mpc': MPC\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition of the setting parameters of the motor and the limit values "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "motor_parameter = dict(r_s=15e-3, l_d=0.37e-3, l_q=1.2e-3, psi_p=65.6e-3, p=3, j_rotor=0.06)\n",
    "limit_values = dict(i=160 * 1.41, omega=12000 * np.pi / 30, u=450)\n",
    "nominal_values = {key: 0.7 * limit for key, limit in limit_values.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition of reference generators for the d- and q- components of the current"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_generator = SwitchedReferenceGenerator(\n",
    "        sub_generators=[\n",
    "            SinusoidalReferenceGenerator(reference_state='i_sq', amplitude_range=(0, 0.3), offset_range=(0, 0.2)),\n",
    "            StepReferenceGenerator(reference_state='i_sq', amplitude_range=(0, 0.5)),\n",
    "            TriangularReferenceGenerator(reference_state='i_sq', amplitude_range=(0, 0.3), offset_range=(0, 0.2))\n",
    "        ],\n",
    "        super_episode_length=(500, 1000)\n",
    "    )\n",
    "\n",
    "d_generator = SwitchedReferenceGenerator(\n",
    "        sub_generators=[\n",
    "            SinusoidalReferenceGenerator(reference_state='i_sd', amplitude_range=(0, 0.3), offset_range=(0, 0.2)),\n",
    "            StepReferenceGenerator(reference_state='i_sd', amplitude_range=(0, 0.5)),\n",
    "            TriangularReferenceGenerator(reference_state='i_sd', amplitude_range=(0, 0.3), offset_range=(0, 0.2))\n",
    "        ],\n",
    "        super_episode_length=(500, 1000)\n",
    "    )\n",
    "\n",
    "reference_generator = MultipleReferenceGenerator([q_generator, d_generator])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, the PMSM is initialized with the previously defined motor parameters and reference generators. The MPC controller is initialized with the specified prediction horizon. Then the motor is simulated step by step. For this purpose the function for visualization is called. The manipulated variables are calculated by the MPC Controller and switched to the PMSM. This takes a long time, because the numerical effort of the MPC is large. Finally, the reward is output.\n",
    "\n",
    "![Graph](img/mpc_currents_voltages.png)\n",
    "The graphs show exemplary the currents and voltages for the control of a reference profile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "#%matplotlib widget\n",
    "# use %matplotlib widget for execution in visual studio code\n",
    "visu = MotorDashboard(state_plots=['i_sq', 'i_sd', 'u_sq', 'u_sd'], update_interval=10) # visualization\n",
    "motor = Motor(MotorType.PermanentMagnetSynchronousMotor,\n",
    "              ControlType.CurrentControl,\n",
    "              ActionType.Continuous)\n",
    "env = gem.make(motor.env_id(),\n",
    "            visualization=visu,\n",
    "            load=ConstantSpeedLoad(omega_fixed=1000 * np.pi / 30), # Fixed preset speed\n",
    "            ode_solver=ScipySolveIvpSolver(), # ODE Solver of the simulation\n",
    "            reference_generator=reference_generator, # initialize the reference generators\n",
    "            \n",
    "            # defenitions for the reward function\n",
    "            reward_function=dict(\n",
    "                reward_weights={'i_sq': 1, 'i_sd': 1},\n",
    "                reward_power=0.5,\n",
    "            ),\n",
    "               \n",
    "            # definitions for the setting parameters\n",
    "            supply=dict(u_nominal=400),\n",
    "            motor=dict(\n",
    "                motor_parameter=motor_parameter,\n",
    "                limit_values=limit_values,\n",
    "                nominal_values=nominal_values\n",
    "            )\n",
    "    )\n",
    "\n",
    "visu.initialize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "controller = Controller.make('mpc', env, ph=3) # initializing the MPC Controller\n",
    "(state, reference), _ = env.reset()\n",
    "cum_rew = 0\n",
    "\n",
    "for i in range(10000): # simulate the PMSM for 10000 steps (1 second)\n",
    "    env.render() # function for visualization\n",
    "    action = controller.control(state, reference) # calculation of the manipulated variables\n",
    "    (state, reference), reward, terminated, truncated, _ = env.step(action) # simulate one step of the PMSM\n",
    "    cum_rew += reward # adding up the Reward \n",
    "    if terminated:\n",
    "        (state, reference), _ = env.reset()\n",
    "        controller.reset()\n",
    "print('Reward =', cum_rew)\n",
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GEM2v0",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
